#!/bin/bash
source "${LMD_BASE_INSTALL_SCRIPT_DIR}/global/sys_check.sh"
source "${LMD_BASE_INSTALL_SCRIPT_DIR}/global/alias.sh"
source "${LMD_BASE_INSTALL_SCRIPT_DIR}/global/conda_run.sh"
source "${LMD_BASE_INSTALL_SCRIPT_DIR}/global/conda_and_pip_settings.sh"
conda config --set show_channel_urls yes

APP_INSTALL_NAME="framepack"
APP_TARGET_DIR_NAME="framepack"
CONDA_ENV_NAME=$APP_INSTALL_NAME
CONDA_PREFIX="./venv"

export CURRENT_APP_PATH="${LMD_APPS_DIR}/${APP_TARGET_DIR_NAME}"
echo 'App Dir: '${CURRENT_APP_PATH}

cd "$CURRENT_APP_PATH"
python --version
pip --version

if [ -z "$HF_ENDPOINT" ] && [ -n "$HF_MIRROR" ]; then
    export HF_ENDPOINT="$HF_MIRROR"
fi


function install_ffmpeg {
  source "${LMD_BASE_INSTALL_SCRIPT_DIR}/global/install_ffmpeg.sh"
}

function check_ffmpeg_and_install_win {
  PROGRAMS=("ffmpeg" "ffprobe")
  missing_programs=0

  for PROGRAM in "${PROGRAMS[@]}"; do
    if ! command -v ${PROGRAM} >/dev/null 2>&1; then
      $PROGRAM -v
      echo "$PROGRAM not found"
      missing_programs=1
    fi
  done

  if [ $missing_programs -eq 1 ]; then
    echo "need install ffmpeg and ffprobe。"
    install_ffmpeg
  else
    echo "ffmpeg and ffprobe exist"
  fi
}

function check_ffmpeg_and_install {
    if [ -f /System/Library/CoreServices/SystemVersion.plist ]; then
        if conda run -p $CONDA_PREFIX --no-capture-output conda list ffmpeg | grep -q "ffmpeg"; then
            echo "found ffmpeg"
        else
            echo "not found ffmpeg. install it."
            conda_run_conda_forge_install ffmpeg
        fi
    elif current_os_is_windows; then
        check_ffmpeg_and_install_win
    fi
}

export GRADIO_TEMP_DIR="$CURRENT_APP_PATH/temp"

if current_os_is_windows; then
  conda_run_pip_install triton-windows transformers==4.53.3
fi

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

if [ -d "$CONDA_PREFIX" ]; then
    echo "use ./venv to start app. "$CONDA_PREFIX

    export MODELSCOPE_HOME="$CURRENT_APP_PATH/tmp_modelscope_home"
    export MODELSCOPE_CACHE="$CURRENT_APP_PATH/tmp_modelscope_cache"

    check_ffmpeg_and_install
    conda_run_pip_install typing_extensions filelock fsspec
    APP_FILE="demo_gradio_f1.py"
    hf_home_by_modelscope="$CURRENT_APP_PATH/tmp_hf_home_by_modelscope"
    target_model_dir="$hf_home_by_modelscope/hub/models--lllyasviel--FramePack_F1_I2V_HY_20250503"
    if [ -d "$target_model_dir" ]; then
      echo "found hf models in $hf_home_by_modelscope"
      export TRANSFORMERS_OFFLINE=1
      export HF_HUB_OFFLINE=1
      export HF_HUB_DISABLE_SYMLINKS=1
      export TRANSFORMERS_CACHE="$hf_home_by_modelscope/hub"
      export HF_HOME="$hf_home_by_modelscope"
      grep -v "HF_HOME" demo_gradio_f1.py > demo_gradio_f1_temp.py
      
      APP_FILE="demo_gradio_f1_temp.py"
      framepack_model_dir="models--lllyasviel--FramePack_F1_I2V_HY_20250503"
      if current_os_is_windows; then
        echo " -- replace with backslashes -- "
        hf_home_by_modelscope_r1=$(echo "$hf_home_by_modelscope" | sed 's/\\/\\\\/g')
        REPLACE_VALUES="s|HunyuanVideoTransformer3DModelPacked.from_pretrained('lllyasviel/FramePack_F1_I2V_HY_20250503'|HunyuanVideoTransformer3DModelPacked.from_pretrained('${hf_home_by_modelscope_r1}/hub/${framepack_model_dir}/snapshots/latest'|g"
        echo "hf_home_by_modelscope_r1: $hf_home_by_modelscope_r1"
        sed -i'' "$REPLACE_VALUES" $APP_FILE
      else
        hf_home_by_modelscope_r1=$(echo "$hf_home_by_modelscope" | sed 's|\\|/|g')
        REPLACE_VALUES="s|HunyuanVideoTransformer3DModelPacked.from_pretrained('lllyasviel/FramePack_F1_I2V_HY_20250503'|HunyuanVideoTransformer3DModelPacked.from_pretrained('${hf_home_by_modelscope_r1}/hub/${framepack_model_dir}/snapshots/latest'|g"
        echo "hf_home_by_modelscope_r1: $hf_home_by_modelscope_r1"
        sed -i '' "$REPLACE_VALUES" $APP_FILE
      fi
    fi
    conda run -p $CONDA_PREFIX --no-capture-output python -Xfaulthandler $APP_FILE --inbrowser
    exit_code=$?
    if [ -d "$target_model_dir" ]; then
      if [ $exit_code -ne 0 ]; then
        echo " exit_code: ${exit_code} "
        if current_os_is_windows; then
          APP_FILE="demo_gradio_f1_temp2.py"
          
          echo " -- replace with slashes -- "
          grep -v "HF_HOME" demo_gradio_f1.py > demo_gradio_f1_temp2.py

          hf_home_by_modelscope_r2=$(echo "$hf_home_by_modelscope" | sed 's|\\|/|g')
          echo "hf_home_by_modelscope_r2: $hf_home_by_modelscope_r2"
          sed -i'' "s|HunyuanVideoTransformer3DModelPacked.from_pretrained('lllyasviel/FramePack_F1_I2V_HY_20250503'|HunyuanVideoTransformer3DModelPacked.from_pretrained('${hf_home_by_modelscope_r2}/hub/${framepack_model_dir}/snapshots/latest'|g" demo_gradio_f1_temp2.py
          
          conda run -p $CONDA_PREFIX --no-capture-output python -Xfaulthandler $APP_FILE --inbrowser
        fi
      fi
    fi
else
    echo "no ./venv. use env. "$CONDA_ENV_NAME

    check_ffmpeg_and_install
    conda run -n $CONDA_ENV_NAME --no-capture-output pip install typing_extensions filelock fsspec
    conda run -n $CONDA_ENV_NAME --no-capture-output python demo_gradio_f1.py --inbrowser
    
fi

